#!/usr/bin/python
# -*- coding: utf-8 -*-
# @version        : 1.0
# @Create Time    : 2021/10/24 16:04 
# @File           : views.py
# @IDE            : PyCharm
# @desc           : 聊天室接口文件


import time
from datetime import datetime
from typing import List
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from utils.response import SuccessResponse

app = APIRouter()


@app.get("/", summary="聊天室接口状态")
def app_root(version: str):
    data = {
        "当前时间": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
        "test": "聊天室",
        "version": version
    }
    return SuccessResponse(data)


class ConnectionManager:
    def __init__(self):
        # 存放激活的ws连接对象
        self.active_connections: List[WebSocket] = []

    async def connect(self, ws: WebSocket):
        # 等待连接
        await ws.accept()
        # 存储ws连接对象
        self.active_connections.append(ws)

    def disconnect(self, ws: WebSocket):
        # 关闭时 移除ws对象
        self.active_connections.remove(ws)

    @staticmethod
    async def send_personal_message(message: str, ws: WebSocket):
        # 发送个人消息
        await ws.send_text(message)

    async def broadcast(self, message: str):
        # 广播消息
        for connection in self.active_connections:
            await connection.send_text(message)


manager = ConnectionManager()


@app.websocket("/chat/{user}")
async def websocket_endpoint(websocket: WebSocket, user: str):

    await manager.connect(websocket)

    await manager.broadcast(f"欢迎-{user}-进入聊天室（{datetime.strftime(datetime.now(), '%Y-%m-%d %H:%M:%S')}）")

    try:
        while True:
            data = await websocket.receive_text()
            # await manager.send_personal_message(f"你说了: {data}", websocket)
            await manager.broadcast(f"{user} 说（{datetime.strftime(datetime.now(), '%Y-%m-%d %H:%M:%S')}）: {data}")

    except WebSocketDisconnect:
        manager.disconnect(websocket)
        await manager.broadcast(f"{user}-退出聊天室")